import socket
import psutil
import threading
import time
import requests
from collections import defaultdict
from datetime import datetime
import json
from concurrent.futures import ThreadPoolExecutor
import queue
import geoip2.database
import os


class NetworkMonitor:
    def __init__(self, geoip_db_path="Country.mmdb"):
        self.connections = []
        self.active_connections = {}
        self.location_cache = {}
        self.running = False
        self.lock = threading.Lock()
        self.monitor_thread = None

        # 添加线程池用于异步处理地理位置查询
        self.executor = ThreadPoolExecutor(max_workers=3)
        self.location_queue = queue.Queue()

        # 初始化GeoIP2数据库
        self.geoip_db_path = geoip_db_path
        self.geoip_reader = None
        self.init_geoip_db()

    def init_geoip_db(self):
        """初始化GeoIP2数据库"""
        try:
            if os.path.exists(self.geoip_db_path):
                self.geoip_reader = geoip2.database.Reader(self.geoip_db_path)
                print(f"GeoIP2数据库加载成功: {self.geoip_db_path}")
            else:
                print(f"警告: GeoIP2数据库文件不存在: {self.geoip_db_path}")
                print("请下载GeoLite2-City.mmdb数据库文件并放置在正确位置")
                print("数据库可从MaxMind官网下载: https://dev.maxmind.com/geoip/geolite2-free-geolocation-data")
        except Exception as e:
            print(f"初始化GeoIP2数据库失败: {e}")
            print("将使用默认位置信息")

    def start(self):
        """启动网络监控"""
        self.running = True
        self.monitor_thread = threading.Thread(target=self._monitor_loop, daemon=True)
        self.monitor_thread.start()

        # 启动地理位置查询线程
        self.location_thread = threading.Thread(target=self._location_worker, daemon=True)
        self.location_thread.start()

        print("网络监控已启动")

    def stop(self):
        """停止网络监控"""
        self.running = False
        self.executor.shutdown(wait=False)

        # 关闭GeoIP2数据库
        if self.geoip_reader:
            self.geoip_reader.close()

        print("网络监控已停止")

    def _monitor_loop(self):
        """网络监控主循环"""
        while self.running:
            try:
                current_connections = psutil.net_connections(kind='inet')
                self._process_connections(current_connections)
                time.sleep(10)  # 增加到5秒检查一次，减少CPU占用
            except Exception as e:
                print(f"网络监控错误: {e}")
                time.sleep(10)

    def _process_connections(self, connections):
        """处理网络连接数据"""
        with self.lock:
            new_connections = []

            for conn in connections:
                if conn.status == 'ESTABLISHED' and conn.raddr:
                    conn_info = {
                        'local_ip': conn.laddr.ip,
                        'local_port': conn.laddr.port,
                        'remote_ip': conn.raddr.ip,
                        'remote_port': conn.raddr.port,
                        'status': conn.status,
                        'pid': conn.pid,
                        'timestamp': time.time()
                    }

                    # 获取进程信息
                    try:
                        if conn.pid:
                            process = psutil.Process(conn.pid)
                            conn_info['process_name'] = process.name()
                        else:
                            conn_info['process_name'] = 'Unknown'
                    except:
                        conn_info['process_name'] = 'Unknown'

                    # 异步获取地理位置信息
                    location_info = self._get_location_info_cached(conn.raddr.ip)
                    conn_info.update(location_info)

                    new_connections.append(conn_info)

            self.connections = new_connections[:]

    def _get_location_info_cached(self, ip):
        """获取IP地址的地理位置信息（带缓存）"""
        if ip in self.location_cache:
            return self.location_cache[ip]

        # 跳过本地IP
        if ip.startswith(('127.', '192.168.', '10.', '172.')):
            location_info = {
                'country': '本地',
                'city': '本地网络',
                'latitude': 39.9042,  # 北京坐标
                'longitude': 116.4074,
                'isp': '本地网络'
            }
            self.location_cache[ip] = location_info
            return location_info

        # 如果不在缓存中，添加到队列异步处理
        if ip not in self.location_cache:
            self.location_queue.put(ip)
            # 返回默认值
            return self._default_location()

        return self.location_cache.get(ip, self._default_location())

    def _location_worker(self):
        """地理位置查询工作线程"""
        while self.running:
            try:
                ip = self.location_queue.get(timeout=1)
                if ip not in self.location_cache:
                    location_info = self._fetch_location_from_geoip(ip)
                    self.location_cache[ip] = location_info
                    print(f"获取IP {ip} 位置信息: {location_info}")
            except queue.Empty:
                continue
            except Exception as e:
                print(f"地理位置查询错误: {e}")

    def _fetch_location_from_geoip(self, ip):
        """通过GeoIP2数据库获取IP地址的地理位置信息"""
        try:
            if not self.geoip_reader:
                return self._fetch_location_backup(ip)

            # 使用GeoIP2查询IP信息
            response = self.geoip_reader.city(ip)

            # 提取地理位置信息
            country = response.country.name or response.country.iso_code
            if not country and response.registered_country:
                country = response.registered_country.name or response.registered_country.iso_code

            city = response.city.name

            # 提取经纬度
            latitude = response.location.latitude if response.location else 0
            longitude = response.location.longitude if response.location else 0

            # 提取ISP信息
            isp = response.traits.isp if response.traits else '未知'

            # 如果是中文环境，尝试获取中文名称
            if hasattr(response.country, 'names') and 'zh-CN' in response.country.names:
                country = response.country.names['zh-CN']

            if hasattr(response.city, 'names') and 'zh-CN' in response.city.names:
                city = response.city.names['zh-CN']

            return {
                'country': country or '未知',
                'city': city or '未知',
                'latitude': latitude or 0,
                'longitude': longitude or 0,
                'isp': isp or '未知'
            }

        except Exception as e:
            print(f"GeoIP2查询失败 {ip}: {e}")
            # 如果GeoIP2查询失败，尝试备用API
            return self._fetch_location_backup(ip)

    def _fetch_location_backup(self, ip):
        """备用IP地理位置查询方法"""
        try:
            # 使用免费的IP地理位置API作为备用
            response = requests.get(f'http://ip-api.com/json/{ip}', timeout=5)
            if response.status_code == 200:
                data = response.json()
                if data['status'] == 'success':
                    return {
                        'country': data.get('country', '未知'),
                        'city': data.get('city', '未知'),
                        'latitude': data.get('lat', 0),
                        'longitude': data.get('lon', 0),
                        'isp': data.get('isp', '未知')
                    }
        except Exception as e:
            print(f"备用API获取地理位置失败 {ip}: {e}")

        return self._default_location()

    def _default_location(self):
        """默认位置信息"""
        return {
            'country': '未知',
            'city': '未知',
            'latitude': 0,
            'longitude': 0,
            'isp': '未知'
        }

    def get_connections(self):
        """获取当前连接列表"""
        with self.lock:
            return self.connections.copy()

    def get_statistics(self):
        """获取网络统计信息"""
        with self.lock:
            total_connections = len(self.connections)
            unique_ips = len(set(conn['remote_ip'] for conn in self.connections))
            countries = list(set(conn.get('country', '未知') for conn in self.connections))

            # 按国家统计连接数
            country_stats = defaultdict(int)
            for conn in self.connections:
                country_stats[conn.get('country', '未知')] += 1

            return {
                'total_connections': total_connections,
                'unique_ips': unique_ips,
                'countries': countries,
                'country_stats': dict(country_stats)
            }
